# -*- coding: utf-8 -*-
import asyncio
import os

import botpy
from botpy import logging
from botpy.ext.cog_yaml import read
from botpy.message import GroupMessage, Message
from botpy.manage import GroupManageEvent

import erniebot
from erniebot_agent.memory import LimitTokensMemory
from erniebot_agent.memory import HumanMessage, AIMessage, SystemMessage, FunctionMessage
from erniebot_agent.chat_models import ERNIEBot
from qianfan.resources.tools import tokenizer
import erniebot.utils

import config
import traceback
from utils import calc_tokens
os.environ["EB_AGENT_ACCESS_TOKEN"] = "212363ee6c08595ab150141603a849f21488c0be"
model = ERNIEBot(model="ernie-4.0-turbo-8k",system=config.SYSTEM)
erniebot.api_type = 'aistudio'
erniebot.access_token = '212363ee6c08595ab150141603a849f21488c0be'
messages = [
]
_log = logging.get_logger()
_log.info("ok!")
class ChatClient(botpy.Client):
    async def on_group_at_message_create(self, message: GroupMessage):
        global messages
        _log.info("收到群聊消息：" + str(message))
        try:
            while 1:
                if messages and isinstance(messages[0], AIMessage):
                    messages = messages[1:]
                    continue
                if messages and isinstance(messages[-1], HumanMessage):
                    messages = messages[:-1]
                    continue
                break
        
            tokens = calc_tokens(messages)
            if tokens > config.MAX_TOKENS:
                messages = messages[2:]
            messages.append(HumanMessage(content=message.content))
            response = await model.chat(messages=messages,system=config.SYSTEM)
            content = response.content
            aimessage = AIMessage(content=content)
            messages.append(aimessage)
            if "作为一个人工智能语言模型" in content:
                messages = messages[:-2]
            if len(content) > config.MAX_MESSAGE_LENGTH:
                # 分割消息为不超过1000字的多个部分
                parts = [content[i:i+config.MAX_MESSAGE_LENGTH] for i in range(0, len(content), config.MAX_MESSAGE_LENGTH)]
            else:
                # 如果不超过1000字，直接作为一个部分发送
                parts = [content]
            # 循环发送每个部分
            for part in parts:
                try:
                    messageResult = await message._api.post_group_message(
                        group_openid=message.group_openid,
                        msg_type=0, 
                        msg_id=message.id,
                        content=part)
                    _log.info(messageResult)
                    await asyncio.sleep(2)  # 在发送每条消息之后等待一秒
                except Exception as e:
                    _log.info(traceback.format_exc())
                    await asyncio.sleep(1)  # 在发送每条消息之后等待一秒

        except:
            _log.info(traceback.format_exc())
            _log.info(messages)
            messages = messages[:-2]
            pass
    async def on_group_add_robot(self, event: GroupManageEvent):
        _log.info("机器人被添加到群聊：" + str(event))
        await self.api.post_group_message(
            group_openid=event.group_openid,
            msg_type=0,
            event_id=event.event_id,
            content="hello",
        )

    async def on_group_del_robot(self, event: GroupManageEvent):
        _log.info("机器人被移除群聊：" + str(event))

    async def on_group_msg_reject(self, event: GroupManageEvent):
        _log.info("群聊关闭机器人主动消息：" + str(event))

    async def on_group_msg_receive(self, event: GroupManageEvent):
        _log.info("群聊打开机器人主动消息：" + str(event))